import torch
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import pdist, squareform
from scipy.cluster.hierarchy import dendrogram, linkage


# 计算簇之间的距离（欧氏距离）
def compute_distance_matrix(X):
    return squareform(pdist(X, 'euclidean'))


# 凝聚型层次聚类
def agglomerative_clustering(X, num_clusters=1):
    # 初始每个点为一个簇
    clusters = [[i] for i in range(X.shape[0])]

    # 计算初始距离矩阵
    dist_matrix = compute_distance_matrix(X)

    # 逐步合并簇直到达到预定的簇数
    while len(clusters) > num_clusters:
        # 找到距离最小的簇对
        min_dist_idx = np.unravel_index(np.argmin(dist_matrix + np.eye(len(dist_matrix)) * np.inf), dist_matrix.shape)
        i, j = min_dist_idx

        # 合并这两个簇
        new_cluster = clusters[i] + clusters[j]

        # 删除旧簇
        if i > j:
            clusters.pop(i)
            clusters.pop(j)
        else:
            clusters.pop(j)
            clusters.pop(i)

        # 更新距离矩阵
        # 计算新簇与所有其他簇的距离
        new_distances = []
        for k in range(len(dist_matrix)):
            if k != i and k != j:
                dist_i = dist_matrix[k][i] if k < i else dist_matrix[k][i - 1]
                dist_j = dist_matrix[k][j] if k < j else dist_matrix[k][j - 1]
                new_distances.append(min(dist_i, dist_j))

        new_distances = np.array(new_distances)

        # 生成新的距离矩阵
        dist_matrix = np.delete(dist_matrix, [i, j], axis=0)
        dist_matrix = np.delete(dist_matrix, [i, j], axis=1)

        # 在 dist_matrix 的末尾添加新簇的距离
        dist_matrix = np.vstack([dist_matrix, new_distances])  # 添加新行
        new_distances = np.append(new_distances, 0)  # 为了和列对齐，添加最后一列
        dist_matrix = np.column_stack([dist_matrix, new_distances])  # 添加新列

        # 打印当前簇的数量
        print(f'当前簇数量：{len(clusters)}')
        print(f'簇结构：{clusters}')

    return clusters


# 生成更多的示例数据：50个二维数据点
np.random.seed(42)
X_np = np.random.rand(50, 2) * 10  # 50个数据点，数据范围在[0, 10]

# 聚类成5个簇
num_clusters = 5
clusters = agglomerative_clustering(X_np, num_clusters=num_clusters)

# 将每个数据点分配到对应的簇
labels = np.zeros(X_np.shape[0])
for idx, cluster in enumerate(clusters):
    for i in cluster:
        labels[i] = idx

# 可视化结果
plt.figure(figsize=(8, 6))

# 按照簇分配颜色
for cluster_idx in range(num_clusters):
    cluster_points = X_np[labels == cluster_idx]  # 使用 NumPy 数组进行索引
    plt.scatter(cluster_points[:, 0], cluster_points[:, 1], label=f"Cluster {cluster_idx + 1}")

plt.title(f'Agglomerative Clustering with {num_clusters} Clusters')
plt.xlabel('X1')
plt.ylabel('X2')
plt.legend()
plt.show()

# 画出层次聚类树状图
linked = linkage(X_np, 'ward')

plt.figure(figsize=(10, 7))
dendrogram(linked)
plt.title('Hierarchical Clustering Dendrogram')
plt.xlabel('Sample Index')
plt.ylabel('Distance')
plt.show()
